#!/usr/bin/env python3
"""
Production deployment test script for Manga Colorizer
Tests complete business workflow without Docker
"""

import os
import sys
import time
import json
import requests
from pathlib import Path
from PIL import Image
import io
import tempfile
import hashlib

# Configuration
API_BASE_URL = "http://localhost:8000"
TEST_TIMEOUT = 30
MAX_RETRIES = 10

class ProductionTester:
    """Complete production testing suite"""
    
    def __init__(self):
        self.test_results = []
        self.passed = 0
        self.total = 0
        
    def log_test(self, test_name: str, passed: bool, message: str = ""):
        """Log test result"""
        self.total += 1
        if passed:
            self.passed += 1
            print(f"✅ {test_name}: {message}")
        else:
            print(f"❌ {test_name}: {message}")
        self.test_results.append({
            "test": test_name,
            "passed": passed,
            "message": message
        })
    
    def create_test_image(self, width=512, height=512) -> bytes:
        """Create a test manga sketch image"""
        img = Image.new('RGB', (width, height), color='white')
        
        # Add manga-style sketch elements
        from PIL import ImageDraw
        draw = ImageDraw.Draw(img)
        
        # Draw a simple manga character face
        # Head outline
        draw.ellipse([(150, 100), (350, 300)], outline='black', width=3)
        
        # Eyes
        draw.ellipse([(200, 150), (220, 170)], outline='black', width=2)
        draw.ellipse([(280, 150), (300, 170)], outline='black', width=2)
        
        # Nose
        draw.line([(250, 180), (250, 200)], fill='black', width=2)
        
        # Mouth
        draw.line([(220, 220), (280, 220)], fill='black', width=2)
        
        # Hair outline
        draw.arc([(140, 90), (360, 110)], 0, 180, fill='black', width=3)
        
        # Convert to bytes
        buffer = io.BytesIO()
        img.save(buffer, format='PNG', quality=95)
        return buffer.getvalue()
    
    def test_server_startup(self) -> bool:
        """Test server startup and health check"""
        try:
            response = requests.get(f"{API_BASE_URL}/api/v1/health", timeout=5)
            if response.status_code == 200:
                data = response.json()
                if data.get('status') == 'healthy':
                    self.log_test("Server Startup", True, "Server is healthy")
                    return True
                else:
                    self.log_test("Server Startup", False, f"Server reports unhealthy: {data}")
                    return False
            else:
                self.log_test("Server Startup", False, f"Status code: {response.status_code}")
                return False
        except requests.exceptions.ConnectionError:
            self.log_test("Server Startup", False, "Cannot connect to server")
            return False
        except Exception as e:
            self.log_test("Server Startup", False, str(e))
            return False
    
    def test_file_upload(self) -> bool:
        """Test file upload functionality"""
        try:
            test_image = self.create_test_image()
            
            files = {'file': ('test_sketch.png', test_image, 'image/png')}
            data = {
                'style': 'anime',
                'strength': 0.7,
                'guidance_scale': 7.5,
                'num_inference_steps': 10
            }
            
            response = requests.post(
                f"{API_BASE_URL}/api/v1/colorize",
                files=files,
                data=data,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                if result.get('success') and 'task_id' in result:
                    self.log_test("File Upload", True, f"Task created: {result['task_id']}")
                    return result['task_id']
                else:
                    self.log_test("File Upload", False, f"Invalid response: {result}")
                    return None
            else:
                self.log_test("File Upload", False, f"Status: {response.status_code}")
                return None
                
        except Exception as e:
            self.log_test("File Upload", False, str(e))
            return None
    
    def test_task_polling(self, task_id: str) -> bool:
        """Test task status polling"""
        try:
            for attempt in range(MAX_RETRIES):
                response = requests.get(
                    f"{API_BASE_URL}/api/v1/colorize/status/{task_id}",
                    timeout=10
                )
                
                if response.status_code == 200:
                    status = response.json()
                    print(f"   Poll {attempt + 1}: {status['status']} ({status.get('progress', 0)}%)")
                    
                    if status['status'] == 'completed':
                        self.log_test("Task Completion", True, f"Completed in {attempt + 1} polls")
                        return True
                    elif status['status'] == 'failed':
                        self.log_test("Task Completion", False, status.get('error_message', 'Unknown error'))
                        return False
                
                time.sleep(3)
            
            self.log_test("Task Completion", False, "Timeout waiting for completion")
            return False
            
        except Exception as e:
            self.log_test("Task Completion", False, str(e))
            return False
    
    def test_result_retrieval(self, task_id: str) -> bool:
        """Test result retrieval"""
        try:
            response = requests.get(
                f"{API_BASE_URL}/api/v1/colorize/result/{task_id}",
                timeout=10
            )
            
            if response.status_code == 200:
                result = response.json()
                if result.get('success') and 'colored_image_url' in result:
                    self.log_test("Result Retrieval", True, "Result retrieved successfully")
                    return True
                else:
                    self.log_test("Result Retrieval", False, "Invalid result format")
                    return False
            else:
                self.log_test("Result Retrieval", False, f"Status: {response.status_code}")
                return False
                
        except Exception as e:
            self.log_test("Result Retrieval", False, str(e))
            return False
    
    def test_batch_processing(self) -> bool:
        """Test batch processing"""
        try:
            # Create multiple test images
            images = []
            for i in range(2):
                image_data = self.create_test_image(256, 256)
                images.append(('files', (f'test_{i}.png', image_data, 'image/png')))
            
            data = {
                'style': 'cartoon',
                'strength': 0.8,
                'guidance_scale': 8.0,
                'num_inference_steps': 8
            }
            
            response = requests.post(
                f"{API_BASE_URL}/api/v1/colorize/batch",
                files=images,
                data=data,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                if 'batch_id' in result and 'task_ids' in result:
                    self.log_test("Batch Processing", True, f"Batch created with {len(result['task_ids'])} tasks")
                    return True
                else:
                    self.log_test("Batch Processing", False, "Invalid batch response")
                    return False
            else:
                self.log_test("Batch Processing", False, f"Status: {response.status_code}")
                return False
                
        except Exception as e:
            self.log_test("Batch Processing", False, str(e))
            return False
    
    def test_error_handling(self) -> bool:
        """Test error handling scenarios"""
        tests = [
            ("Invalid file type", {'file': ('test.txt', b'invalid', 'text/plain')}),
            ("Empty file", {'file': ('empty.png', b'', 'image/png')}),
            ("Invalid style", {'file': ('test.png', self.create_test_image(), 'image/png'), 'style': 'invalid'}),
            ("Large file", {'file': ('large.png', b'x' * 11000000, 'image/png')}),  # 11MB
        ]
        
        all_passed = True
        for test_name, data in tests:
            try:
                response = requests.post(
                    f"{API_BASE_URL}/api/v1/colorize",
                    files=data if 'file' in data else None,
                    data={k: v for k, v in data.items() if k != 'file'},
                    timeout=10
                )
                
                if response.status_code in [400, 413, 422]:
                    self.log_test(f"Error Handling - {test_name}", True, f"Properly rejected: {response.status_code}")
                else:
                    self.log_test(f"Error Handling - {test_name}", False, f"Unexpected status: {response.status_code}")
                    all_passed = False
                    
            except Exception as e:
                self.log_test(f"Error Handling - {test_name}", False, str(e))
                all_passed = False
        
        return all_passed
    
    def test_api_endpoints(self) -> bool:
        """Test all API endpoints"""
        endpoints = [
            ("GET", f"{API_BASE_URL}/api/v1/health"),
            ("GET", f"{API_BASE_URL}/api/v1/colorize/tasks"),
            ("GET", f"{API_BASE_URL}/docs"),
        ]
        
        all_passed = True
        for method, url in endpoints:
            try:
                if method == "GET":
                    response = requests.get(url, timeout=5)
                    if response.status_code == 200:
                        self.log_test(f"API Endpoint - {url}", True, "Accessible")
                    else:
                        self.log_test(f"API Endpoint - {url}", False, f"Status: {response.status_code}")
                        all_passed = False
            except Exception as e:
                self.log_test(f"API Endpoint - {url}", False, str(e))
                all_passed = False
        
        return all_passed
    
    def run_complete_test_suite(self) -> dict:
        """Run complete production test suite"""
        print("🧪 Manga Colorizer - Production Deployment Test")
        print("=" * 60)
        
        # Test 1: Server startup
        print("\n1. Testing server startup...")
        if not self.test_server_startup():
            return {"success": False, "message": "Server not running"}
        
        # Test 2: API endpoints
        print("\n2. Testing API endpoints...")
        self.test_api_endpoints()
        
        # Test 3: File upload
        print("\n3. Testing file upload...")
        task_id = self.test_file_upload()
        
        if task_id:
            # Test 4: Task processing
            print("\n4. Testing task processing...")
            self.test_task_polling(task_id)
            
            # Test 5: Result retrieval
            print("\n5. Testing result retrieval...")
            self.test_result_retrieval(task_id)
        
        # Test 6: Batch processing
        print("\n6. Testing batch processing...")
        self.test_batch_processing()
        
        # Test 7: Error handling
        print("\n7. Testing error handling...")
        self.test_error_handling()
        
        # Summary
        print("\n" + "=" * 60)
        print(f"🎯 Production Test Results: {self.passed}/{self.total} tests passed")
        
        success = self.passed == self.total
        if success:
            print("🎉 All tests passed! System is production ready.")
        else:
            print("❌ Some tests failed. Check the logs above.")
        
        return {
            "success": success,
            "passed": self.passed,
            "total": self.total,
            "results": self.test_results
        }


def create_client_demo():
    """Create client demo script"""
    demo_script = '''#!/usr/bin/env python3
"""
Client demo script for Manga Colorizer API
"""

import requests
import time
from PIL import Image
import io

API_BASE_URL = "http://localhost:8000/api/v1"

def demo_colorization():
    """Demo the complete colorization workflow"""
    print("🎨 Manga Colorizer Client Demo")
    print("=" * 40)
    
    # Create test image
    img = Image.new('RGB', (512, 512), color='white')
    from PIL import ImageDraw
    draw = ImageDraw.Draw(img)
    draw.ellipse([(150, 100), (350, 300)], outline='black', width=3)
    
    buffer = io.BytesIO()
    img.save(buffer, format='PNG')
    image_data = buffer.getvalue()
    
    # Step 1: Upload image
    print("1. Uploading image...")
    files = {'file': ('demo_sketch.png', image_data, 'image/png')}
    data = {'style': 'anime', 'strength': 0.7}
    
    response = requests.post(f"{API_BASE_URL}/colorize", files=files, data=data)
    if response.status_code != 200:
        print(f"❌ Upload failed: {response.text}")
        return
    
    result = response.json()
    task_id = result['task_id']
    print(f"✅ Task created: {task_id}")
    
    # Step 2: Poll for completion
    print("2. Processing...")
    for i in range(30):
        status = requests.get(f"{API_BASE_URL}/colorize/status/{task_id}").json()
        print(f"   Status: {status['status']} ({status.get('progress', 0)}%)")
        
        if status['status'] == 'completed':
            break
        elif status['status'] == 'failed':
            print(f"❌ Processing failed: {status.get('error_message')}")
            return
        
        time.sleep(2)
    
    # Step 3: Get result
    print("3. Getting result...")
    result = requests.get(f"{API_BASE_URL}/colorize/result/{task_id}").json()
    print(f"✅ Colorization complete!")
    print(f"   Original: {result.get('original_image_url')}")
    print(f"   Colored: {result.get('colored_image_url')}")
    print(f"   Processing time: {result.get('processing_time', 0):.2f}s")

if __name__ == "__main__":
    demo_colorization()
'''
    
    with open('scripts/client_demo.py', 'w') as f:
        f.write(demo_script)
    
    print("✅ Client demo script created: scripts/client_demo.py")


if __name__ == "__main__":
    tester = ProductionTester()
    results = tester.run_complete_test_suite()
    
    # Create client demo
    create_client_demo()
    
    # Exit with appropriate code
    sys.exit(0 if results["success"] else 1)
